#include <stdio.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <arpa/inet.h>
#include <unistd.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/stat.h>
#include <fcntl.h>
#include <string.h>
#include <errno.h>
#include <openssl/sha.h> 
// 需要安装OpenSSL库: sudo apt-get install libssl-dev
#define N 64
#define BUFFER_SIZE 1024
#define MAX_FILENAME_SIZE 256

// 计算文件的SHA256哈希值
int calculate_sha256(const char *filename, unsigned char *hash) {
    FILE *file = fopen(filename, "rb");
    if (!file) {
        perror("打开文件失败");
        return -1;
    }
    
    SHA256_CTX sha256;
    SHA256_Init(&sha256);
    
    unsigned char buffer[BUFFER_SIZE];
    size_t bytes_read;
    
    while ((bytes_read = fread(buffer, 1, BUFFER_SIZE, file))) {
        SHA256_Update(&sha256, buffer, bytes_read);
    }
    
    SHA256_Final(hash, &sha256);
    fclose(file);
    
    return 0;
}

// 接收文件函数
int receive_file(int sockfd, char *received_filename) {
    char file_info[MAX_FILENAME_SIZE + 20];
    memset(file_info, 0, sizeof(file_info));
    
    // 接收文件信息（文件名和文件大小）
    ssize_t info_bytes = recv(sockfd, file_info, sizeof(file_info) - 1, 0); // 保留一个字节给字符串结束符
    if (info_bytes <= 0) {
        if (info_bytes == 0) {
            fprintf(stderr, "连接被对端关闭\n");
        } else {
            perror("接收文件信息失败");
        }
        return -1;
    }
    
    // 确保字符串以null结尾
    file_info[info_bytes] = '\0';
    printf("接收到的文件信息: %s\n", file_info);
    printf("接收到的信息长度: %zd\n", info_bytes);
    
    // 解析文件名和文件大小
    char *token = strtok(file_info, "|");
    if (token == NULL) {
        fprintf(stderr, "文件信息格式错误: %s\n", file_info);
        return -1;
    }
    
    char filename[MAX_FILENAME_SIZE];
    strncpy(filename, token, sizeof(filename) - 1);
    filename[sizeof(filename) - 1] = '\0';
    
    // 复制文件名到输出参数
    strncpy(received_filename, filename, MAX_FILENAME_SIZE - 1);
    received_filename[MAX_FILENAME_SIZE - 1] = '\0';
    
    token = strtok(NULL, "|");
    if (token == NULL) {
        fprintf(stderr, "文件信息格式错误，缺少文件大小\n");
        return -1;
    }
    
    long file_size = atol(token);
    printf("接收升级文件: %s, 大小: %ld 字节\n", filename, file_size);
    
    // 接收文件哈希
    unsigned char expected_hash[SHA256_DIGEST_LENGTH];
    ssize_t hash_bytes = recv(sockfd, expected_hash, sizeof(expected_hash), MSG_WAITALL);
    if (hash_bytes != sizeof(expected_hash)) {
        fprintf(stderr, "接收哈希值失败，期望 %zu 字节，实际 %zd 字节\n", 
                sizeof(expected_hash), hash_bytes);
        return -1;
    }
    
    // 打印接收到的哈希值（用于调试）
    printf("接收到的哈希值: ");
    for (int i = 0; i < SHA256_DIGEST_LENGTH; i++) {
        printf("%02x", expected_hash[i]);
    }
    printf("\n");
    
    // 创建文件 - 使用接收到的文件名，保存在当前目录
    FILE *fp = fopen(filename, "wb");
    if (fp == NULL) {
        perror("创建文件失败");
        return -1;
    }
    
    // 接收文件内容
    char buffer[BUFFER_SIZE];
    long total_received = 0;
    ssize_t bytes_received;
    
    while (total_received < file_size) {
        bytes_received = recv(sockfd, buffer, BUFFER_SIZE, 0);
        if (bytes_received <= 0) {
            if (bytes_received == 0) {
                printf("连接已关闭\n");
            } else {
                perror("接收数据失败");
            }
            fclose(fp);
            return -1;
        }
        
        // 写入文件
        size_t bytes_written = fwrite(buffer, 1, bytes_received, fp);
        if (bytes_written != bytes_received) {
            perror("写入文件失败");
            fclose(fp);
            return -1;
        }
        
        total_received += bytes_received;
        
        // 显示进度
        printf("已接收: %.2f%%\r", (float)total_received / file_size * 100);
        fflush(stdout);
    }
    
    printf("\n文件接收完成，共接收 %ld 字节\n", total_received);
    fclose(fp);
    
    // 验证文件大小
    if (total_received != file_size) {
        fprintf(stderr, "文件大小不匹配: 期望 %ld, 实际 %ld\n", file_size, total_received);
        return -1;
    }
    
    // 验证文件完整性
    unsigned char actual_hash[SHA256_DIGEST_LENGTH];
    if (calculate_sha256(filename, actual_hash) != 0) {
        fprintf(stderr, "计算文件哈希失败\n");
        return -1;
    }
    
    // 打印计算出的哈希值（用于调试）
    printf("计算出的哈希值: ");
    for (int i = 0; i < SHA256_DIGEST_LENGTH; i++) {
        printf("%02x", actual_hash[i]);
    }
    printf("\n");
    
    if (memcmp(expected_hash, actual_hash, SHA256_DIGEST_LENGTH) != 0) {
        fprintf(stderr, "文件完整性验证失败\n");
        return -1;
    }
    
    printf("文件完整性验证成功\n");
    
    // 发送接收结果确认
    const char *result = "SUCCESS";
    if (send(sockfd, result, strlen(result), 0) < 0) {
        perror("发送确认失败");
        return -1;
    }
    
    return 0;
}

// 执行升级程序
int perform_upgrade(const char *filename) {
    printf("开始执行升级程序: %s\n", filename);
    
    // 这里应该是实际的升级逻辑
    // 例如: 验证固件签名、写入闪存、重启设备等
    
    // 模拟升级过程
    printf("验证固件签名...\n");
    sleep(1);
    printf("备份当前系统...\n");
    sleep(1);
    printf("写入闪存...\n");
    sleep(2);
    printf("升级完成! 需要重启设备以应用更新\n");
    
    return 0;
}

int main(int argc, char const *argv[])
{
    int broadfd;

    //创建一个socket文件描述符
    broadfd = socket(AF_INET, SOCK_DGRAM, 0);
    if(broadfd < 0)
    {
        perror("sock err");
        return -1;
    }
    //设置套接字选项，允许广播
    int broadcast = 1;
    if (setsockopt(broadfd, SOL_SOCKET, SO_BROADCAST, &broadcast, sizeof(broadcast)) < 0) {
        perror("setsockopt broadcast err");
        close(broadfd);
        return -1;
    }

    //绑定套接字(ip+port)
    struct sockaddr_in addr;
    addr.sin_family = AF_INET;
    addr.sin_port = htons(7878);  //端口号
    addr.sin_addr.s_addr = INADDR_ANY;

    int addrlen = sizeof(addr);
    
    if(bind(broadfd, (struct sockaddr *)&addr, addrlen) < 0)
    {
        perror("bind err");
        return -1;
    }

    ssize_t len;
    char buf[N] = {0};
    struct sockaddr_in cliaddr;
    
    //接收"jinan"搜索包
    bzero(buf, N);
    len = recvfrom(broadfd, buf, N, 0, (struct sockaddr *)&cliaddr, &addrlen);
    printf("接收到广播: %s\n", buf);
    
    //判断是否是本公司产品：收到的数据是否"jinan"
    if(strcmp(buf, "jinan") != 0)
    {
        printf("不是本公司的设备\n");
        close(broadfd);
        return -1;
    }

    printf("收到搜索请求，回复确认\n");
    
    //回复yes，告诉软件，我收到了搜索协议，并且回复地址
    if(sendto(broadfd, "yes", 4, 0, (struct sockaddr *)&cliaddr, addrlen) < 0)
    {
        perror("sendto err");
        close(broadfd);
        return -1;
    }

    //关闭广播套接字
    close(broadfd);
        
    //变身为TCP服务器，准备接收软件的升级文件
    int tcpfd = socket(AF_INET, SOCK_STREAM, 0);
    if(tcpfd < 0)
    {
        perror("sock err");
        return -1;
    }

    //设置SO_REUSEADDR选项，避免地址占用错误
    int reuse = 1;
    if (setsockopt(tcpfd, SOL_SOCKET, SO_REUSEADDR, &reuse, sizeof(reuse)) < 0) {
        perror("setsockopt reuse err");
        close(tcpfd);
        return -1;
    }

    if(bind(tcpfd, (struct sockaddr *)&addr, addrlen) < 0)
    {
        perror("bind err");
        close(tcpfd);
        return -1;
    }

    //监听套接字
    if(listen(tcpfd, 5) < 0)
    {
        perror("listen err");
        close(tcpfd);
        return -1;
    }

    //接收客户端的连接
    printf("等待客户端连接...\n");
    int clifd;
    
    //接收对端的地址
    struct sockaddr_in client_addr;
    socklen_t client_len = sizeof(client_addr);
    clifd = accept(tcpfd, (struct sockaddr *)&client_addr, &client_len);
    if(clifd < 0)
    {
        perror("accept err");
        close(tcpfd);
        return -1;
    }

    printf("新的连接建立，客户端IP: %s, 端口: %d\n", 
           inet_ntoa(client_addr.sin_addr), ntohs(client_addr.sin_port));
    printf("准备接收升级文件\n");    

    // 接收升级文件
    char received_filename[MAX_FILENAME_SIZE];
    if (receive_file(clifd, received_filename) == 0) {
        printf("升级文件接收成功\n");
        
        // 使用实际接收到的文件名进行升级
        if (perform_upgrade(received_filename) == 0) {
            printf("升级过程完成\n");
        } else {
            printf("升级过程失败\n");
        }
    } else {
        printf("升级文件接收失败\n");
    }

    close(clifd);
    close(tcpfd);
    
    return 0;
}